package cn.oatn;

import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

public class T1008_ASM_AmendMethod {

    public static void test() throws IOException {
        ClassReader cr = new ClassReader("cn.oatn.MyRunnable");
        ClassWriter cw = new ClassWriter(0);

        // 注意cw作为第二个参数
        cr.accept(new ClassVisitor(Opcodes.ASM9, cw) {
            @Override
            public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
                MethodVisitor originalMV = super.visitMethod(access, name, descriptor, signature, exceptions);
                if (name.equals("a")) { // 简单判断
                    return new MethodVisitor(Opcodes.ASM9, originalMV) {
                        // 变量序号，暂定100，不能和已有的冲突
                        int startTimeVarIndex = 100;

                        @Override
                        public void visitCode() {
                            super.visitCode();
                            // 开始之后插入一段代码记录局部变量
                            visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J", false);
                            visitVarInsn(Opcodes.LSTORE, startTimeVarIndex);
                        }

                        @Override
                        public void visitInsn(int opcode) {
                            if (opcode == Opcodes.IRETURN) {
                                visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
                                // return之前插入一段代码计算耗时
                                visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J", false);
                                visitVarInsn(Opcodes.LLOAD, startTimeVarIndex);
                                visitInsn(Opcodes.LSUB);
                                visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(J)V", false);
                            }
                            super.visitInsn(opcode);
                        }
                    };
                } else {
                    return originalMV;
                }
            }
        }, ClassReader.EXPAND_FRAMES);
        OutputStream o = new FileOutputStream("XXX.class");
        o.write(cw.toByteArray());
        o.close();
    }
}
